%% JRD model parameter recovery

% First, lets recover parameters for full reference direction users

numParticipants = 14;

% Initial parameter values
x0 = [.4 4];

% Use this for global search
opts = optimset('Algorithm','interior-point');
problem = createOptimProblem('fmincon','options',opts,'objective',@JRDLogGrowthModelFullRefFixL_ParamRecovery, ... 
    'x0',x0,'lb',[0,0]);
gs = GlobalSearch;

% Create a matrix for reference directions
refDirection = [0,1; 1,0; 0,-1; -1,0];
global refDirectionFull
refDirectionFull = 1;

headingConditions = [0,45,90,135,180,225,270,315];
pointConditions = headingConditions;

for p=1:numParticipants
    
    % Parameter: noisy object location
    a(p,1) = optFull(p,1);
    b(p,1) = optFull(p,2);
    l = 1;
    q = 20;

    % Create a matrix containing the locations of all objects relative to an
    % origin (obj id is the index)
    locRelOrigin = [-1,0; 0,-1; 0,1; 1,0; 1,1; 1,2; 2,-1; 2,0; 2,1];

    % Create cell array; each index is an item containing a n by 2 matrix of 
    % vectors representing the distances of each item from that location
    % relative to the reference axes. The encoded object locations are then
    % computed from those.
    for i=1:length(locRelOrigin(:,1))
        for j=1:length(locRelOrigin(:,1))
            locRelObj{i}(j,1) = locRelOrigin(j,1) - locRelOrigin(i,1);
            locRelObj{i}(j,2) = locRelOrigin(j,2) - locRelOrigin(i,2);
        end
    end

    % Now create an array of test probes (first index = standing location;
    % second index = imagined heading; third index = pointing direction) using
    % combination algorithm where order matters
    n = length(locRelOrigin(:,1)); k = 3;
    nk = nchoosek(1:n,k);
    testProbe=zeros(0,k);
    for i=1:size(nk,1)
        pi = perms(nk(i,:));
        testProbe = unique([testProbe; pi],'rows');
    end

    % Organize test probes by heading and allocentric pointing direction (round
    % to get rid of strange decimals)
    for i=1:length(testProbe(:,1))
        if thetaFunc([1,0],locRelObj{testProbe(i,1)}(testProbe(i,2),:)) <= ... 
                thetaFunc([-1,0],locRelObj{testProbe(i,1)}(testProbe(i,2),:))
            testProbe(i,4) = round(thetaFunc([0,1],locRelObj{testProbe(i,1)}(testProbe(i,2),:)));
        else
            testProbe(i,4) = round(180 + thetaFunc([0,-1],locRelObj{testProbe(i,1)}(testProbe(i,2),:)));
        end
        if thetaFunc([1,0],locRelObj{testProbe(i,1)}(testProbe(i,3),:)) <= ... 
                thetaFunc([-1,0],locRelObj{testProbe(i,1)}(testProbe(i,3),:))
            testProbe(i,5) = round(thetaFunc([0,1],locRelObj{testProbe(i,1)}(testProbe(i,3),:)));
        else
            testProbe(i,5) = round(180 + thetaFunc([0,-1],locRelObj{testProbe(i,1)}(testProbe(i,3),:)));
        end
    end

    % Need to duplicate trials for heading and allocentric condition that only
    % appear once given the arrangement of objects
    for i=1:length(testProbe(:,1))
        if (testProbe(i,4) == 45 && testProbe(i,5) == 180) || (testProbe(i,4) == ...
                180 && testProbe(i,5) == 45) || (testProbe(i,4) == 45 && ...
                testProbe(i,5) == 270) || (testProbe(i,4) == 270 && testProbe(i,5) ...
                == 45) || (testProbe(i,4) == 135 && testProbe(i,5) == 315) || ...
                (testProbe(i,4) == 315 && testProbe(i,5) == 135)
            testProbe(length(testProbe(:,1)) + 1,:) = testProbe(i,:);
        end
    end
    
     % Randomize sample of test probes
    testProbe = testProbe(randperm(size(testProbe,1)),:);
    pointResponse_recov = [];
    
    global newTestProbe
    % Get 128 test probes, a trial for each heading and pointing direction
    count = zeros(8,8);
    newProbeCount = 0;
    for i=1:length(testProbe(:,1))
        for j=1:length(headingConditions)
            if round(testProbe(i,4)) == round(headingConditions(j)) && sum(count(j,:)) < 16
                for k=1:length(pointConditions)
                    if round(testProbe(i,5)) == round(pointConditions(k)) && count(j,k) < 2
                        count(j,k) = count(j,k) + 1;
                        newProbeCount = newProbeCount + 1;
                        newTestProbe(newProbeCount,:) = testProbe(i,:);
                    end
                end
            end
        end
    end
    
    % Randomize new test probes
    newTestProbe = newTestProbe(randperm(size(newTestProbe,1)),:);
            
    % Objects with noisy encoding
    for i=1:length(locRelOrigin(:,1))
        for j=1:length(locRelOrigin(:,1))
            for k=1:length(refDirection(:,1))
                angleHolder(k) = round(thetaFunc(refDirection(k,:),locRelObj{i}(j,:)));           
            end
            [M,I] = min(angleHolder);
            locEncoded{i}(j,1) = locRelObj{i}(j,1) + normrnd(0,(a(p,1)/(1+b(p,1)*exp(-(l)*thetaFunc(refDirection(I,:),locRelObj{i}(j,:))))));
            locEncoded{i}(j,2) = locRelObj{i}(j,2) + normrnd(0,(a(p,1)/(1+b(p,1)*exp(-(l)*thetaFunc(refDirection(I,:),locRelObj{i}(j,:))))));
            if i > 1
                while i - j > 0    
                    locEncoded{i}(i-j,:) = -(locEncoded{i-j}(i,:));
                    j = j + 1;
                end
            end
        end
    end

    % Now an algorithm to cycle through the test probes and perform the JRD
    % task
    for i=1:length(newTestProbe(:,1))

        % Define standing position, imagined heading, and pointing direction
        standpos = newTestProbe(i,1);
        imagheading = newTestProbe(i,2);
        pointobj = newTestProbe(i,3);

        % Create arrays for heading condition and correct pointing direction
        if thetaFunc([1,0],locRelObj{standpos}(imagheading,:)) <= ... 
                thetaFunc([-1,0],locRelObj{standpos}(imagheading,:))        
            headingCond(i) = thetaFunc([0,1],locRelObj{standpos}(imagheading,:));
        else
            headingCond(i) = 180 + thetaFunc([0,-1],locRelObj{standpos}(imagheading,:));
        end
        % Get rid of strange decimals
        headingCond(i) = round(headingCond(i));
        
        correctPointAngle(i) = atanThetaFunc(locRelObj{standpos}(imagheading,:), ...
              locRelObj{standpos}(pointobj,:));
        % Get rid of strange decimals
        correctPointAngle(i) = round(correctPointAngle(i));

        % Create array for pointing condition
        if thetaFunc([1,0],locRelObj{standpos}(pointobj,:)) <= ... 
                thetaFunc([-1,0],locRelObj{standpos}(pointobj,:))        
            pointingCond(i) = thetaFunc([0,1],locRelObj{standpos}(pointobj,:));
        else
            pointingCond(i) = 180 + thetaFunc([0,-1],locRelObj{standpos}(pointobj,:));
        end
        % Get rid of strange decimals
        pointingCond(i) = round(pointingCond(i));
        
        % Check if imagined heading is consistent with reference direction
        for j=1:length(refDirection(:,1))
            if round(thetaFunc(locEncoded{standpos}(imagheading,:),refDirection(j,:))) <= q
                pointResponse_recov(i) = atanThetaFunc(refDirection(j,:), ... 
                    locEncoded{standpos}(pointobj,:));
                break
            end
        end
        
        % If imagined direction is not in line, need to do additional
        % computaions (vecter addition)
        if length(pointResponse_recov) ~= i
            angleVector = locEncoded{standpos}(imagheading,:) + ... 
                locEncoded{imagheading}(pointobj,:);
            pointResponse_recov(i) = atanThetaFunc(locEncoded{standpos}(imagheading,:), ... 
                angleVector);
        end  
    end
    
    global pointingError_recov
    for i=1:length(pointResponse_recov)
        pointingError_recov(i) = abs(pointResponse_recov(i) - correctPointAngle(i));
        if pointingError_recov(i) > 180
            pointingError_recov(i) = 360 - pointingError_recov(i);
        end
    end
    
    % Run optimization
    optFull_recov(p,:) = run(gs,problem);

end

%% Now for the 0-180 reference users

% Create a matrix for reference directions
refDirection = [0,1; 0,-1;];
global refDirectionFull
refDirectionFull = 0;

for p=1:numParticipants
    
    % Parameter: noisy object location
    a(p,1) = optHalf(p,1);
    b(p,1) = optHalf(p,2);
    l = 1;
    q = 20;

    % Create a matrix containing the locations of all objects relative to an
    % origin (obj id is the index)
    locRelOrigin = [-1,0; 0,-1; 0,1; 1,0; 1,1; 1,2; 2,-1; 2,0; 2,1];

    % Create cell array; each index is an item containing a n by 2 matrix of 
    % vectors representing the distances of each item from that location
    % relative to the reference axes. The encoded object locations are then
    % computed from those.
    for i=1:length(locRelOrigin(:,1))
        for j=1:length(locRelOrigin(:,1))
            locRelObj{i}(j,1) = locRelOrigin(j,1) - locRelOrigin(i,1);
            locRelObj{i}(j,2) = locRelOrigin(j,2) - locRelOrigin(i,2);
        end
    end

    % Now create an array of test probes (first index = standing location;
    % second index = imagined heading; third index = pointing direction) using
    % combination algorithm where order matters
    n = length(locRelOrigin(:,1)); k = 3;
    nk = nchoosek(1:n,k);
    testProbe=zeros(0,k);
    for i=1:size(nk,1)
        pi = perms(nk(i,:));
        testProbe = unique([testProbe; pi],'rows');
    end

    % Organize test probes by heading and allocentric pointing direction (round
    % to get rid of strange decimals)
    for i=1:length(testProbe(:,1))
        if thetaFunc([1,0],locRelObj{testProbe(i,1)}(testProbe(i,2),:)) <= ... 
                thetaFunc([-1,0],locRelObj{testProbe(i,1)}(testProbe(i,2),:))
            testProbe(i,4) = round(thetaFunc([0,1],locRelObj{testProbe(i,1)}(testProbe(i,2),:)));
        else
            testProbe(i,4) = round(180 + thetaFunc([0,-1],locRelObj{testProbe(i,1)}(testProbe(i,2),:)));
        end
        if thetaFunc([1,0],locRelObj{testProbe(i,1)}(testProbe(i,3),:)) <= ... 
                thetaFunc([-1,0],locRelObj{testProbe(i,1)}(testProbe(i,3),:))
            testProbe(i,5) = round(thetaFunc([0,1],locRelObj{testProbe(i,1)}(testProbe(i,3),:)));
        else
            testProbe(i,5) = round(180 + thetaFunc([0,-1],locRelObj{testProbe(i,1)}(testProbe(i,3),:)));
        end
    end

    % Need to duplicate trials for heading and allocentric condition that only
    % appear once given the arrangement of objects
    for i=1:length(testProbe(:,1))
        if (testProbe(i,4) == 45 && testProbe(i,5) == 180) || (testProbe(i,4) == ...
                180 && testProbe(i,5) == 45) || (testProbe(i,4) == 45 && ...
                testProbe(i,5) == 270) || (testProbe(i,4) == 270 && testProbe(i,5) ...
                == 45) || (testProbe(i,4) == 135 && testProbe(i,5) == 315) || ...
                (testProbe(i,4) == 315 && testProbe(i,5) == 135)
            testProbe(length(testProbe(:,1)) + 1,:) = testProbe(i,:);
        end
    end
    
     % Randomize sample of test probes
    testProbe = testProbe(randperm(size(testProbe,1)),:);
    pointResponse_recov = [];
    
    global newTestProbe
    % Get 128 test probes, a trial for each heading and pointing direction
    count = zeros(8,8);
    newProbeCount = 0;
    for i=1:length(testProbe(:,1))
        for j=1:length(headingConditions)
            if round(testProbe(i,4)) == round(headingConditions(j)) && sum(count(j,:)) < 16
                for k=1:length(pointConditions)
                    if round(testProbe(i,5)) == round(pointConditions(k)) && count(j,k) < 2
                        count(j,k) = count(j,k) + 1;
                        newProbeCount = newProbeCount + 1;
                        newTestProbe(newProbeCount,:) = testProbe(i,:);
                    end
                end
            end
        end
    end
    
    % Randomize new test probes
    newTestProbe = newTestProbe(randperm(size(newTestProbe,1)),:);
            
    % Objects with noisy encoding
    for i=1:length(locRelOrigin(:,1))
        for j=1:length(locRelOrigin(:,1))
            for k=1:length(refDirection(:,1))
                angleHolder(k) = round(thetaFunc(refDirection(k,:),locRelObj{i}(j,:)));           
            end
            [M,I] = min(angleHolder);
            locEncoded{i}(j,1) = locRelObj{i}(j,1) + normrnd(0,(a(p,1)/(1+b(p,1)*exp(-(l)*thetaFunc(refDirection(I,:),locRelObj{i}(j,:))))));
            locEncoded{i}(j,2) = locRelObj{i}(j,2) + normrnd(0,(a(p,1)/(1+b(p,1)*exp(-(l)*thetaFunc(refDirection(I,:),locRelObj{i}(j,:))))));
            if i > 1
                while i - j > 0    
                    locEncoded{i}(i-j,:) = -(locEncoded{i-j}(i,:));
                    j = j + 1;
                end
            end
        end
    end

    % Now an algorithm to cycle through the test probes and perform the JRD
    % task
    for i=1:length(newTestProbe(:,1))

        % Define standing position, imagined heading, and pointing direction
        standpos = newTestProbe(i,1);
        imagheading = newTestProbe(i,2);
        pointobj = newTestProbe(i,3);

        % Create arrays for heading condition and correct pointing direction
        if thetaFunc([1,0],locRelObj{standpos}(imagheading,:)) <= ... 
                thetaFunc([-1,0],locRelObj{standpos}(imagheading,:))        
            headingCond(i) = thetaFunc([0,1],locRelObj{standpos}(imagheading,:));
        else
            headingCond(i) = 180 + thetaFunc([0,-1],locRelObj{standpos}(imagheading,:));
        end
        % Get rid of strange decimals
        headingCond(i) = round(headingCond(i));
        
        correctPointAngle(i) = atanThetaFunc(locRelObj{standpos}(imagheading,:), ...
              locRelObj{standpos}(pointobj,:));
        % Get rid of strange decimals
        correctPointAngle(i) = round(correctPointAngle(i));

        % Create array for pointing condition
        if thetaFunc([1,0],locRelObj{standpos}(pointobj,:)) <= ... 
                thetaFunc([-1,0],locRelObj{standpos}(pointobj,:))        
            pointingCond(i) = thetaFunc([0,1],locRelObj{standpos}(pointobj,:));
        else
            pointingCond(i) = 180 + thetaFunc([0,-1],locRelObj{standpos}(pointobj,:));
        end
        % Get rid of strange decimals
        pointingCond(i) = round(pointingCond(i));
        
        % Check if imagined heading is consistent with reference direction
        for j=1:length(refDirection(:,1))
            if round(thetaFunc(locEncoded{standpos}(imagheading,:),refDirection(j,:))) <= q
                pointResponse_recov(i) = atanThetaFunc(refDirection(j,:), ... 
                    locEncoded{standpos}(pointobj,:));
                break
            end
        end
        
        % If imagined direction is not in line, need to do additional
        % computaions (vecter addition)
        if length(pointResponse_recov) ~= i
            angleVector = locEncoded{standpos}(imagheading,:) + ... 
                locEncoded{imagheading}(pointobj,:);
            pointResponse_recov(i) = atanThetaFunc(locEncoded{standpos}(imagheading,:), ... 
                angleVector);
        end  
    end
    
    global pointingError_recov
    for i=1:length(pointResponse_recov)
        pointingError_recov(i) = abs(pointResponse_recov(i) - correctPointAngle(i));
        if pointingError_recov(i) > 180
            pointingError_recov(i) = 360 - pointingError_recov(i);
        end
    end
    
    % Run optimization
    optHalf_recov(p,:) = run(gs,problem);

end

%% Plot results
corr_a_Full = corrcoef(optFull(:,1),optFull_recov(:,1));
figure;
scatter(optFull(:,1),optFull_recov(:,1));

corr_b_Full = corrcoef(optFull(:,2),optFull_recov(:,2));
figure;
scatter(optFull(:,2),optFull_recov(:,2));

corr_a_Half = corrcoef(optHalf(:,1),optHalf_recov(:,1));
figure;
scatter(optHalf(:,1),optHalf_recov(:,1));

corr_b_Half = corrcoef(optHalf(:,2),optHalf_recov(:,2));
figure;
scatter(optHalf(:,2),optHalf_recov(:,2));